"""
@Description :   GLM-4.5V 模型 (api)

@Author      :   tqychy 
@Time        :   2025/08/23 22:16:40
"""
import sys

sys.path.append("./nets")

import base64

from api_key import api_key
from zai import ZhipuAiClient


class GLM_4_5_V:
    def __init__(self, *args, model_path):
        self.cfg, self.logger = args
        self.model_name = model_path
        self.model = ZhipuAiClient(api_key=api_key("zhipuai"))

    def __call__(self, inputs):
        response = self.model.chat.completions.create(
            model=self.model_name,
            messages=inputs,
            thinking={
                "type": "disabled"
            }
        )

        return response.choices[0].message.content
    
    @staticmethod
    def make_prompt(image_path, sentence, category, mode):
        if mode == "rec":
            prompt = f"""
            <image>
            Here is a description of the objects in the figure. 
            Please enclose the corresponding positions using coordinate boxes. 
            Examples of coordinate value formats: [x1, y1, x2, y2]
            Please only output boxes in a fixed format like:
            ```json\n[\n\t{{"bbox_2d": [x1, y1, x2, y2], "label": "..."}}\n]\n```
            and do not output any unnecessary content.

            Description:\n{sentence}
            """
        elif mode == "detect":
            prompt = f"""
            <image>
            Here is a series of category labels separated by "." and an image. Please identify ALL objects in the image corresponding to the mentioned category labels, and meet the following requirements:
            1. There may be multiple objects for each category label; please identify ALL of them.
            2. Return a JSON-formatted list of dictionaries. Each dictionary must contain the "bbox", "category" and "score" of the identified object. An example of the output format is: [{{"bbox": [456, 249, 507, 324], "category": "Dog", "score": 0.67}}, {{"bbox": [495, 261, 528, 322], "category": "Dog", "score": 0.84}}, {{"bbox": [392, 323, 447, 401], "category": "Monkey", "score": 0.46}}].
            3. Do not output any extra information.
            4. The returned category labels must EXACTLY match the input category labels.
            Categories: {category}
            """
        else:
            raise ValueError(f"Unknown task type! {mode}")

        with open(image_path, "rb") as img_file:
            img_base = base64.b64encode(img_file.read()).decode("utf-8")
        inputs = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": img_base
                        }
                    },
                    {
                        "type": "text",
                        "text": prompt
                    }
                ]
            }
        ]
        return inputs
    
    @staticmethod
    def covert_formatted_bbox(bbox, image_shape):
        """
        将归一化并缩放的边界框格式 [x1, y1, x2, y2] 转换为 [x_min, y_min, width, height] 格式。
        """
        height, width, _ = image_shape
        x1, y1, x2, y2 = bbox
        # 反归一化
        x_min = (x1 / 1000.0) * width
        y_min = (y1 / 1000.0) * height
        x_max = (x2 / 1000.0) * width
        y_max = (y2 / 1000.0) * height
        # 计算宽度和高度
        width_abs = x_max - x_min
        height_abs = y_max - y_min
        return [x_min, y_min, width_abs, height_abs]


if __name__ == "__main__":
    image_path = "./dataset/scripts/refcoco/data/images/mscoco/images/train2014/COCO_train2014_000000000009.jpg"
    with open(image_path, "rb") as img_file:
        img_base = base64.b64encode(img_file.read()).decode("utf-8")
    inputs = [
        {
            "role": "user",
            "content": [
                {
                    "type": "image_url",
                    "image_url": {
                        "url": img_base
                    }
                },
                {
                    "type": "text",
                    "text": "请描述这个图片"
                }
            ]
        }
    ]

    model = GLM_4_5_V(None, None)
    print(model(inputs))




        